#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Population First-Order PBPK Modeling Script

Simulates PBPK modeling for a population of 100 individuals using first-order kinetics
with population variability in key parameters.

Author: PBPK Modeling System
Version: 1.0.0
Date: 2025-07-24
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import json
import os
from typing import Dict, Any, List
import warnings
warnings.filterwarnings('ignore')

# Handle np.trapz deprecation with fallback compatibility
try:
    from scipy.integrate import trapezoid
    def trapz_func(y, x=None, dx=1.0, axis=-1):
        """Fallback-compatible trapezoidal integration using scipy.integrate.trapezoid"""
        return trapezoid(y, x=x, dx=dx, axis=axis)
except ImportError:
    try:
        from scipy.integrate import trapz as trapz_func
    except ImportError:
        # Fallback to numpy's trapz if scipy is not available
        trapz_func = np.trapz

class PopulationFirstOrderPBPK:
    """Population-based first-order PBPK model"""
    
    def __init__(self):
        """Initialize population model"""
        self.population_size = 100
        self.base_params = {
            'drug_name': 'Generic_Drug',
            'dose': 100.0,  # mg
            'ka': 1.0,      # 1/h absorption rate
            'ke': 0.15,     # 1/h elimination rate
            'Vd': 50.0,     # L distribution volume
            'F': 0.85,      # bioavailability
            'tlag': 0.1,    # h lag time
            'weight': 70.0, # kg
            'age': 35.0,    # years
            'gender': 'mixed'
        }
        self.population = []
        self.simulation_results = {}
        
    def generate_population(self, population_type: str = "healthy_adults") -> List[Dict[str, float]]:
        """Generate population with realistic variability"""
        np.random.seed(42)  # For reproducibility
        
        if population_type == "healthy_adults":
            # Normal distribution parameters based on literature
            weight_mean, weight_std = 70.0, 10.0
            age_mean, age_std = 35.0, 12.0
            
            # PK parameter variability (CV%)
            cv_ka = 0.25  # 25% CV for absorption
            cv_ke = 0.30  # 30% CV for elimination
            cv_Vd = 0.20  # 20% CV for volume
            cv_F = 0.15   # 15% CV for bioavailability
            
            population = []
            for i in range(self.population_size):
                # Generate individual characteristics
                weight = np.random.normal(weight_mean, weight_std)
                weight = max(45, min(120, weight))  # Clamp to realistic range
                
                age = np.random.normal(age_mean, age_std)
                age = max(18, min(80, age))  # Clamp to adult range
                
                # Generate PK parameters with variability
                individual = {
                    'id': f'individual_{i+1:03d}',
                    'weight': weight,
                    'age': age,
                    'gender': np.random.choice(['male', 'female'], p=[0.5, 0.5]),
                    'dose': self.base_params['dose'],
                    'ka': np.random.lognormal(np.log(self.base_params['ka']), cv_ka),
                    'ke': np.random.lognormal(np.log(self.base_params['ke']), cv_ke),
                    'Vd': np.random.lognormal(np.log(self.base_params['Vd'] * weight/70), cv_Vd),
                    'F': np.random.beta(8.5, 1.5) * 0.9 + 0.1,  # Beta distribution for F
                    'tlag': max(0, np.random.normal(self.base_params['tlag'], 0.05))
                }
                
                # Ensure parameter bounds
                individual['ka'] = max(0.1, min(5.0, individual['ka']))
                individual['ke'] = max(0.01, min(1.0, individual['ke']))
                individual['Vd'] = max(10, min(200, individual['Vd']))
                individual['F'] = max(0.3, min(1.0, individual['F']))
                
                population.append(individual)
            
            return population
    
    def first_order_absorption(self, individual: Dict[str, float], time_points: np.ndarray) -> Dict[str, Any]:
        """Calculate first-order absorption for an individual"""
        dose = individual['dose'] * individual['F']
        ka = individual['ka']
        tlag = individual['tlag']
        
        # Time with lag
        t_lag = np.maximum(0, time_points - tlag)
        
        # Amount absorbed over time
        absorbed_amount = dose * (1 - np.exp(-ka * t_lag))
        absorbed_amount[time_points < tlag] = 0
        
        # Absorption rate
        absorption_rate = np.zeros_like(time_points)
        mask = time_points > tlag
        absorption_rate[mask] = dose * ka * np.exp(-ka * (time_points[mask] - tlag))
        
        return {
            'absorbed_amount': absorbed_amount,
            'absorption_rate': absorption_rate,
            'total_absorbed': absorbed_amount[-1] if len(absorbed_amount) > 0 else 0
        }
    
    def one_compartment_model(self, individual: Dict[str, float], time_points: np.ndarray) -> Dict[str, Any]:
        """Calculate one-compartment PK for an individual"""
        dose = individual['dose'] * individual['F']
        ke = individual['ke']
        Vd = individual['Vd']
        ka = individual['ka']
        tlag = individual['tlag']
        
        # Time with lag
        t_lag = np.maximum(0, time_points - tlag)
        
        # Concentration = (F * Dose * ka) / (Vd * (ka - ke)) * (e^(-ke*t) - e^(-ka*t))
        if abs(ka - ke) < 1e-6:  # Handle ka == ke case
            concentration = (dose * ka * t_lag * np.exp(-ke * t_lag)) / Vd
        else:
            concentration = (dose * ka) / (Vd * (ka - ke)) * (np.exp(-ke * t_lag) - np.exp(-ka * t_lag))
        
        concentration[time_points < tlag] = 0
        
        # Calculate key metrics
        if len(concentration) > 0:
            cmax = np.max(concentration)
            tmax_idx = np.argmax(concentration)
            tmax = time_points[tmax_idx]
            auc = trapz_func(concentration, time_points)
            
            # Terminal half-life
            terminal_points = concentration[-min(10, len(concentration)//3):]
            if len(terminal_points) > 1 and np.any(terminal_points > 0):
                log_conc = np.log(terminal_points[terminal_points > 0])
                time_term = time_points[-len(terminal_points):][terminal_points > 0]
                if len(log_conc) > 1:
                    slope, _ = np.polyfit(time_term, log_conc, 1)
                    half_life = -0.693 / slope
                else:
                    half_life = 0.693 / ke
            else:
                half_life = 0.693 / ke
        else:
            cmax = tmax = auc = half_life = 0
        
        return {
            'concentration': concentration,
            'cmax': cmax,
            'tmax': tmax,
            'auc': auc,
            'half_life': half_life,
            'clearance': dose / auc if auc > 0 else 0
        }
    
    def run_population_simulation(self, simulation_time: float = 24.0, time_step: float = 0.1) -> Dict[str, Any]:
        """Run simulation for entire population"""
        print(f"Running population simulation for {self.population_size} individuals...")
        
        # Generate population
        population = self.generate_population()
        
        # Create time points
        time_points = np.arange(0, simulation_time + time_step, time_step)
        
        # Store individual results
        individual_results = []
        population_concentrations = []
        
        for i, individual in enumerate(population):
            # Calculate absorption and PK
            absorption = self.first_order_absorption(individual, time_points)
            pk = self.one_compartment_model(individual, time_points)
            
            # Store results
            individual_result = {
                'id': individual['id'],
                'demographics': {
                    'weight': individual['weight'],
                    'age': individual['age'],
                    'gender': individual['gender']
                },
                'parameters': {
                    'ka': individual['ka'],
                    'ke': individual['ke'],
                    'Vd': individual['Vd'],
                    'F': individual['F'],
                    'tlag': individual['tlag']
                },
                'pk_metrics': {
                    'cmax': pk['cmax'],
                    'tmax': pk['tmax'],
                    'auc': pk['auc'],
                    'half_life': pk['half_life'],
                    'clearance': pk['clearance']
                },
                'time_points': time_points,
                'concentration': pk['concentration']
            }
            
            individual_results.append(individual_result)
            population_concentrations.append(pk['concentration'])
        
        # Calculate population statistics
        concentrations_array = np.array(population_concentrations)
        
        population_stats = {
            'mean_concentration': np.mean(concentrations_array, axis=0),
            'std_concentration': np.std(concentrations_array, axis=0),
            'median_concentration': np.median(concentrations_array, axis=0),
            'percentile_5': np.percentile(concentrations_array, 5, axis=0),
            'percentile_95': np.percentile(concentrations_array, 95, axis=0),
            'time_points': time_points
        }
        
        # Calculate population PK metrics
        cmax_values = [ind['pk_metrics']['cmax'] for ind in individual_results]
        auc_values = [ind['pk_metrics']['auc'] for ind in individual_results]
        tmax_values = [ind['pk_metrics']['tmax'] for ind in individual_results]
        
        population_pk = {
            'cmax_stats': {
                'mean': np.mean(cmax_values),
                'std': np.std(cmax_values),
                'median': np.median(cmax_values),
                'cv': np.std(cmax_values) / np.mean(cmax_values) * 100 if np.mean(cmax_values) > 0 else 0
            },
            'auc_stats': {
                'mean': np.mean(auc_values),
                'std': np.std(auc_values),
                'median': np.median(auc_values),
                'cv': np.std(auc_values) / np.mean(auc_values) * 100 if np.mean(auc_values) > 0 else 0
            },
            'tmax_stats': {
                'mean': np.mean(tmax_values),
                'std': np.std(tmax_values),
                'median': np.median(tmax_values),
                'cv': np.std(tmax_values) / np.mean(tmax_values) * 100 if np.mean(tmax_values) > 0 else 0
            }
        }
        
        return {
            'population': population,
            'individual_results': individual_results,
            'population_stats': population_stats,
            'population_pk': population_pk,
            'simulation_time': simulation_time,
            'time_step': time_step
        }
    
    def generate_population_charts(self, results: Dict[str, Any], save_path: str = "../output"):
        """Generate comprehensive population visualizations"""
        os.makedirs(save_path, exist_ok=True)
        
        time_points = results['population_stats']['time_points']
        mean_conc = results['population_stats']['mean_concentration']
        std_conc = results['population_stats']['std_concentration']
        p5 = results['population_stats']['percentile_5']
        p95 = results['population_stats']['percentile_95']
        
        # Create comprehensive visualization
        fig, axes = plt.subplots(3, 3, figsize=(24, 18))
        fig.suptitle('Population First-Order PBPK Modeling Results (100 Individuals)', fontsize=16, fontweight='bold')
        
        # 1. Population concentration-time profile
        ax1 = axes[0, 0]
        ax1.plot(time_points, mean_conc, 'b-', linewidth=3, label='Mean')
        ax1.fill_between(time_points, p5, p95, alpha=0.3, color='blue', label='90% CI')
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Plasma Concentration (mg/L)')
        ax1.set_title('Population Concentration-Time Profile')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # 2. Individual trajectories (subset)
        ax2 = axes[0, 1]
        individual_results = results['individual_results']
        for i in range(0, len(individual_results), 5):  # Show every 5th individual
            ax2.plot(time_points, individual_results[i]['concentration'], alpha=0.7, linewidth=1)
        ax2.plot(time_points, mean_conc, 'r-', linewidth=2, label='Mean')
        ax2.set_xlabel('Time (hours)')
        ax2.set_ylabel('Plasma Concentration (mg/L)')
        ax2.set_title('Individual Trajectories (Sample)')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        
        # 3. Cmax distribution
        ax3 = axes[0, 2]
        cmax_values = [ind['pk_metrics']['cmax'] for ind in individual_results]
        ax3.hist(cmax_values, bins=20, alpha=0.7, color='skyblue', edgecolor='black')
        ax3.axvline(np.mean(cmax_values), color='red', linestyle='--', linewidth=2, label=f'Mean: {np.mean(cmax_values):.1f}')
        ax3.set_xlabel('Cmax (mg/L)')
        ax3.set_ylabel('Frequency')
        ax3.set_title('Cmax Distribution')
        ax3.legend()
        
        # 4. AUC distribution
        ax4 = axes[1, 0]
        auc_values = [ind['pk_metrics']['auc'] for ind in individual_results]
        ax4.hist(auc_values, bins=20, alpha=0.7, color='lightgreen', edgecolor='black')
        ax4.axvline(np.mean(auc_values), color='red', linestyle='--', linewidth=2, label=f'Mean: {np.mean(auc_values):.1f}')
        ax4.set_xlabel('AUC (mg·h/L)')
        ax4.set_ylabel('Frequency')
        ax4.set_title('AUC Distribution')
        ax4.legend()
        
        # 5. Parameter correlation: Cmax vs ke
        ax5 = axes[1, 1]
        ke_values = [ind['parameters']['ke'] for ind in individual_results]
        ax5.scatter(ke_values, cmax_values, alpha=0.6, color='purple')
        ax5.set_xlabel('Elimination Rate (1/h)')
        ax5.set_ylabel('Cmax (mg/L)')
        ax5.set_title('Cmax vs Elimination Rate')
        ax5.grid(True, alpha=0.3)
        
        # 6. Parameter correlation: AUC vs Vd
        ax6 = axes[1, 2]
        vd_values = [ind['parameters']['Vd'] for ind in individual_results]
        ax6.scatter(vd_values, auc_values, alpha=0.6, color='orange')
        ax6.set_xlabel('Volume of Distribution (L)')
        ax6.set_ylabel('AUC (mg·h/L)')
        ax6.set_title('AUC vs Vd')
        ax6.grid(True, alpha=0.3)
        
        # 7. Gender comparison
        ax7 = axes[2, 0]
        male_cmax = [ind['pk_metrics']['cmax'] for ind in individual_results if ind['demographics']['gender'] == 'male']
        female_cmax = [ind['pk_metrics']['cmax'] for ind in individual_results if ind['demographics']['gender'] == 'female']
        
        ax7.boxplot([male_cmax, female_cmax], labels=['Male', 'Female'])
        ax7.set_ylabel('Cmax (mg/L)')
        ax7.set_title('Gender Comparison - Cmax')
        
        # 8. Age vs PK parameters
        ax8 = axes[2, 1]
        age_values = [ind['demographics']['age'] for ind in individual_results]
        ax8.scatter(age_values, auc_values, alpha=0.6, color='green')
        z = np.polyfit(age_values, auc_values, 1)
        p = np.poly1d(z)
        ax8.plot(age_values, p(age_values), "r--", alpha=0.8)
        ax8.set_xlabel('Age (years)')
        ax8.set_ylabel('AUC (mg·h/L)')
        ax8.set_title('Age vs AUC')
        ax8.grid(True, alpha=0.3)
        
        # 9. Population summary statistics
        ax9 = axes[2, 2]
        ax9.axis('off')
        
        pk_stats = results['population_pk']
        summary_text = f"""Population Summary (n=100)

Key PK Parameters:
• Cmax: {pk_stats['cmax_stats']['mean']:.1f} ± {pk_stats['cmax_stats']['std']:.1f} mg/L
  (CV: {pk_stats['cmax_stats']['cv']:.1f}%)

• AUC: {pk_stats['auc_stats']['mean']:.1f} ± {pk_stats['auc_stats']['std']:.1f} mg·h/L
  (CV: {pk_stats['auc_stats']['cv']:.1f}%)

• Tmax: {pk_stats['tmax_stats']['mean']:.2f} ± {pk_stats['tmax_stats']['std']:.2f} h

Population Characteristics:
• Age: {np.mean([ind['demographics']['age'] for ind in individual_results]):.1f} ± {np.std([ind['demographics']['age'] for ind in individual_results]):.1f} years
• Weight: {np.mean([ind['demographics']['weight'] for ind in individual_results]):.1f} ± {np.std([ind['demographics']['weight'] for ind in individual_results]):.1f} kg
• Gender: {len([ind for ind in individual_results if ind['demographics']['gender'] == 'male'])}M / {len([ind for ind in individual_results if ind['demographics']['gender'] == 'female'])}F
"""
        
        ax9.text(0.1, 0.9, summary_text, transform=ax9.transAxes, 
                fontsize=11, verticalalignment='top', 
                bbox=dict(boxstyle="round,pad=0.5", facecolor="lightyellow", alpha=0.8))
        
        plt.tight_layout()
        chart_path = os.path.join(save_path, "population_first_order_pbpk.png")
        plt.savefig(chart_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        return chart_path
    
    def export_population_data(self, results: Dict[str, Any], save_path: str = "../output"):
        """Export population data to CSV for further analysis"""
        os.makedirs(save_path, exist_ok=True)
        
        # Create DataFrame with individual results
        individual_results = results['individual_results']
        
        data_list = []
        for ind in individual_results:
            row = {
                'id': ind['id'],
                'weight': ind['demographics']['weight'],
                'age': ind['demographics']['age'],
                'gender': ind['demographics']['gender'],
                'ka': ind['parameters']['ka'],
                'ke': ind['parameters']['ke'],
                'Vd': ind['parameters']['Vd'],
                'F': ind['parameters']['F'],
                'tlag': ind['parameters']['tlag'],
                'cmax': ind['pk_metrics']['cmax'],
                'tmax': ind['pk_metrics']['tmax'],
                'auc': ind['pk_metrics']['auc'],
                'half_life': ind['pk_metrics']['half_life'],
                'clearance': ind['pk_metrics']['clearance']
            }
            data_list.append(row)
        
        df = pd.DataFrame(data_list)
        csv_path = os.path.join(save_path, "population_first_order_results.csv")
        df.to_csv(csv_path, index=False)
        
        print(f"📊 Population data exported to: {csv_path}")
        return csv_path

def main():
    """Main execution function"""
    print("=" * 80)
    print("POPULATION FIRST-ORDER PBPK MODELING")
    print("Simulating 100 individuals with first-order kinetics")
    print("=" * 80)
    
    # Initialize model
    model = PopulationFirstOrderPBPK()
    
    # Run population simulation
    print("\n1. Running population simulation...")
    results = model.run_population_simulation(simulation_time=24.0, time_step=0.1)
    
    # Display population statistics
    print("\n2. Population PK Statistics:")
    print("-" * 40)
    pk_stats = results['population_pk']
    
    for param, stats in pk_stats.items():
        print(f"{param.upper()}:")
        print(f"  Mean: {stats['mean']:.2f}")
        print(f"  Std: {stats['std']:.2f}")
        if 'cv' in stats:
            print(f"  CV%: {stats['cv']:.1f}%")
        print()
    
    # Generate comprehensive visualizations
    print("3. Generating population charts...")
    chart_path = model.generate_population_charts(results)
    print(f"📊 Charts saved to: {chart_path}")
    
    # Export population data
    csv_path = model.export_population_data(results)
    
    # Display sample individual results
    print("\n4. Sample Individual Results:")
    print("-" * 40)
    sample_size = min(5, len(results['individual_results']))
    for i in range(sample_size):
        ind = results['individual_results'][i]
        print(f"Individual {i+1}:")
        print(f"  Age: {ind['demographics']['age']:.1f} years, Weight: {ind['demographics']['weight']:.1f} kg")
        print(f"  Cmax: {ind['pk_metrics']['cmax']:.2f} mg/L, AUC: {ind['pk_metrics']['auc']:.1f} mg·h/L")
        print()
    
    print("=" * 80)
    print("✅ Population First-Order PBPK Modeling Complete!")
    print("=" * 80)

if __name__ == "__main__":
    main()